package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"mcp-db-gateway/pkg/logger"
	"net/http"

	"mcp-db-gateway/pkg/config"
	"mcp-db-gateway/pkg/db"
	"mcp-db-gateway/pkg/storage"
	"mcp-db-gateway/pkg/tools"

	_ "github.com/go-sql-driver/mysql"
	_ "github.com/lib/pq"
	"go.uber.org/zap"
)

// --- Mock MCP Server ---
// This is a simplified simulation of an MCP server.
type ToolFunc func(ctx context.Context, params string) (interface{}, error)

type MCPServer struct {
	tools map[string]ToolFunc
}

func NewMCPServer() *MCPServer {
	return &MCPServer{
		tools: make(map[string]ToolFunc),
	}
}

func (s *MCPServer) RegisterTool(name string, handler ToolFunc) {
	s.tools[name] = handler
	logger.Log.Info("Tool registered", zap.String("tool_name", name))
}

func (s *MCPServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	toolName := r.URL.Path[1:] // Get tool name from URL path
	handler, ok := s.tools[toolName]
	if !ok {
		http.Error(w, "tool not found", http.StatusNotFound)
		return
	}

	if r.Method != http.MethodPost {
		http.Error(w, "invalid method", http.StatusMethodNotAllowed)
		return
	}

	body, err := io.ReadAll(r.Body)
	if err != nil {
		http.Error(w, "failed to read request body", http.StatusBadRequest)
		return
	}
	defer r.Body.Close()

	logger.Log.Debug("Received request",
		zap.String("tool_name", toolName),
		zap.String("remote_addr", r.RemoteAddr),
		zap.String("body", string(body)),
	)

	result, err := handler(r.Context(), string(body))
	if err != nil {
		http.Error(w, fmt.Sprintf("tool execution error: %v", err), http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	if result != nil {
		jsonData, err := json.MarshalIndent(result, "", "  ") // Indent with 2 spaces
		if err != nil {
			http.Error(w, "failed to marshal result to JSON", http.StatusInternalServerError)
			return
		}
		w.Write(jsonData)
	}
}

func (s *MCPServer) Start(addr string) {
	logger.Log.Info("MCP Server starting", zap.String("address", addr))
	if err := http.ListenAndServe(addr, s); err != nil {
		logger.Log.Fatal("MCP Server failed", zap.Error(err))
	}
}

// --- Main Application ---
func main() {
	// 1. Initialize Logger
	logger.InitLogger()
	defer logger.Log.Sync() // Flushes buffer, if any

	// 2. Load general configuration (like server port)
	cfg, err := config.LoadConfig("config.yaml")
	if err != nil {
		logger.Log.Fatal("Failed to load config", zap.Error(err))
	}

	// 3. Initialize storage
	store, err := storage.NewStore("gateway.db")
	if err != nil {
		logger.Log.Fatal("Failed to initialize storage", zap.Error(err))
	}
	logger.Log.Info("Storage initialized successfully")

	// 4. Initialize database manager from storage
	dbManager, err := db.NewManager(store)
	if err != nil {
		logger.Log.Fatal("Failed to initialize db manager", zap.Error(err))
	}
	defer dbManager.Close()
	logger.Log.Info("Database manager initialized successfully")

	// 5. Initialize tools
	dbExecutor := tools.NewDBExecutor(dbManager)
	dbManagerTools := tools.NewDBManagerTools(dbManager)

	// 6. Initialize and start MCP server
	mcpServer := NewMCPServer()

	// 7. Register all tools with the server
	mcpServer.RegisterTool("execute_sql", dbExecutor.ExecuteSQL)
	mcpServer.RegisterTool("add_database", dbManagerTools.AddDatabase)
	mcpServer.RegisterTool("remove_database", dbManagerTools.RemoveDatabase)
	mcpServer.RegisterTool("list_databases", dbManagerTools.ListDatabases)
	mcpServer.RegisterTool("health_check", dbManagerTools.HealthCheck)

	logger.Log.Info("All tools registered. Starting server.")
	addr := fmt.Sprintf(":%d", cfg.Server.Port)
	mcpServer.Start(addr)
} 